

package ai.moran.demo.common.es;

import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.google.gson.Gson;
import com.google.gson.JsonArray;
import com.google.gson.JsonObject;
import io.searchbox.client.JestResult;
import io.searchbox.core.SearchScroll;
import org.apache.commons.beanutils.BeanUtils;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;



public class ScrollSearchResult<T> implements Iterable<List<T>> {
    private boolean isSuccess;
    private String nextScrollId;
    private List<T> items;
    private Class<T> sourceClass;
    private Long total;
    private EsJestService esJestService;
    private String scrollTime;
    private String indexName;
    private String typeName;
    private String idPropertyName;

    private ScrollSearchResult(EsJestService esJestService, String scrollTime, String indexName, String typeName, Class<T> sourceClass, String idPropertyName, boolean isSuccess, String nextScrollId, List<T> items, Long total) {
        this.esJestService = esJestService;
        this.scrollTime = scrollTime;
        this.indexName = indexName;
        this.typeName = typeName;
        this.sourceClass = sourceClass;
        this.idPropertyName = idPropertyName;
        this.isSuccess = isSuccess;
        this.nextScrollId = nextScrollId;
        this.items = items;
        this.total = total;
    }

    public Long getTotal() {
        return this.total;
    }

    public Iterator<List<T>> iterator() {
        return new ScrollSearchResult.ScrollSearchIterator();
    }

    public class ScrollSearchIterator implements Iterator<List<T>> {
        private List<T> nextItems;

        public ScrollSearchIterator() {
        }

        public boolean hasNext() {
            return StringUtils.isNotEmpty(ScrollSearchResult.this.nextScrollId) && ScrollSearchResult.this.isSuccess;
        }

        public List<T> next() {
            try {
                SearchScroll scroll = (new io.searchbox.core.SearchScroll.Builder(ScrollSearchResult.this.nextScrollId, ScrollSearchResult.this.scrollTime)).build();
                JestResult jestResult = ScrollSearchResult.this.esJestService.getJestClient().execute(scroll);
                if (!ScrollSearchResult.this.esJestService.jestResultSucceeded(jestResult, ScrollSearchResult.this.indexName, ScrollSearchResult.this.typeName, (String)null)) {
                    throw new RuntimeException(jestResult.getErrorMessage());
                } else {
                    JsonArray hits = jestResult.getJsonObject().getAsJsonObject("hits").getAsJsonArray("hits");
                    this.nextItems = new ArrayList(hits.size());

                    for(int i = 0; i < hits.size(); ++i) {
                        JsonObject hitJson = hits.get(i).getAsJsonObject();
                        Gson gson = new Gson();
                        T sourceObj = gson.fromJson(hitJson.get("_source").toString(), ScrollSearchResult.this.sourceClass);
                        if (StringUtils.isNotEmpty(ScrollSearchResult.this.idPropertyName)) {
                            BeanUtils.setProperty(sourceObj, ScrollSearchResult.this.idPropertyName, hitJson.get("_id").toString());
                        }

                        this.nextItems.add(sourceObj);
                    }

                    List<T> returnData = new ArrayList(ScrollSearchResult.this.items);
                    ScrollSearchResult.this.items = this.nextItems;
                    if (ScrollSearchResult.this.items.size() == 0) {
                        ScrollSearchResult.this.nextScrollId = "";
                    } else {
                        ScrollSearchResult.this.nextScrollId = jestResult.getJsonObject().get("_scroll_id").getAsString();
                    }

                    return returnData;
                }
            } catch (Exception var8) {
                throw new NoSuchElementException(var8.getMessage());
            }
        }
    }

    public static class Builder<T> {
        private boolean isSuccess;
        private String nextScrollId;
        private List<T> items;
        private Class<T> sourceClass;
        private String idPropertyName;
        private Long total;
        private EsJestService esJestService;
        private String scrollTime;
        private String indexName;
        private String typeName;

        public Builder(EsJestService esJestService, String scrollTime, String indexName, String typeName, Class<T> sourceClass) {
            this.esJestService = esJestService;
            this.scrollTime = scrollTime;
            this.indexName = indexName;
            this.typeName = typeName;
            this.sourceClass = sourceClass;
            this.isSuccess = false;
            this.nextScrollId = "";
            this.items = new ArrayList(0);
        }

        public ScrollSearchResult.Builder<T> isSuccess(boolean isSuccess) {
            this.isSuccess = isSuccess;
            return this;
        }

        public ScrollSearchResult.Builder<T> nextScrollId(String nextScrollId) {
            this.nextScrollId = nextScrollId;
            return this;
        }

        public ScrollSearchResult.Builder<T> dataItems(List<T> items) {
            this.items = items;
            return this;
        }

        public ScrollSearchResult.Builder<T> idPropertyName(String idPropertyName) {
            this.idPropertyName = idPropertyName;
            return this;
        }

        public ScrollSearchResult.Builder<T> total(Long total) {
            this.total = total;
            return this;
        }

        public ScrollSearchResult<T> build() {
            return new ScrollSearchResult(this.esJestService, this.scrollTime, this.indexName, this.typeName, this.sourceClass, this.idPropertyName, this.isSuccess, this.nextScrollId, this.items, this.total);
        }
    }
}

